// Copyright 2022 jeff.li. and/or its affiliates.

#include <../src/codegen/codegen_c_host.h>
#include <gtest/gtest.h>
#include <tbir/ir/expr.h>
#include <tbir/ir/function.h>
#include <tbir/ir/module.h>
#include <tbir/runtime/container.h>
#include <tbir/runtime/registry.h>
#include <iostream>

namespace tbir::ir {
    using namespace runtime;

    TEST(IR, AllocaVarStmt) {
        DataType int_ty = DataType::Int(64);
        const auto *printer = ::tbir::runtime::FunctionRegistry::Get("ir.AsText");
        const auto *build_module = ::tbir::runtime::FunctionRegistry::Get("module.build.c");

        AllocaVarStmt alloca_stmt("b", PrimType(int_ty), IntImm(int_ty, 0));
        AssignStmt assign_stmt(alloca_stmt->var, PrimExpr(10));
        ReturnStmt rt_stmt(alloca_stmt->var);

        Array<Stmt> seq_stmt;
        seq_stmt.push_back(alloca_stmt);
        seq_stmt.push_back(assign_stmt);
        seq_stmt.push_back(rt_stmt);
        SeqStmt body(seq_stmt);
        Array<PrimVar> params{};
        PrimFunc func(params, {}, body, PrimType(DataType::Int(32)));

        String ir_text = (*printer)({func}).As<String>();
        std::cout << ir_text << std::endl;

        func = WithAttr(std::move(func), attr::kGlobalSymbol, StringRef("test_alloca"));

        codegen::CodeGenCHost cg;
        cg.AddFunction(func);
        std::string code = cg.Finish();
        std::cout << code << std::endl;

        IRModule module = IRModule::FromExpr(func);
        ::tbir::runtime::Module m = (*build_module)({module}).As<Module>();
        std::cout << m->GetSource() << std::endl;
    }

}  // namespace tbir::ir
